"""
MCP 服务器模块
提供 Model Context Protocol 服务器实现，集成网页内容提取功能
"""
import os
import logging
import threading
from typing import Optional, Dict, Any

from mcp.server.fastmcp import FastMCP
from tools.web_search import WebSearchTool

# 配置日志
logging.basicConfig(
    format='[%(process)d][%(thread)d][%(filename)s#%(lineno)d] %(levelname)s: %(message)s',
    level=logging.DEBUG
)
logger = logging.getLogger(__name__)

class WebSearchMCP:
    """WebSearch MCP 服务器实现"""
    
    def __init__(self, name: str = "WebSearch MCP", output_dir: str = "./results", 
                 headless: bool = True, format: str = "markdown", extract_mode: str = "main",
                 file: Optional[str] = None):
        """
        初始化 WebSearch MCP 服务器
        
        参数:
            name: 服务器名称
            output_dir: 搜索结果输出目录
            headless: 是否使用无头浏览器模式
            format: 输出格式 (markdown 或 text)
            extract_mode: 内容提取模式 (main 或 body)
            file: 包含URL列表的文件路径
        """
        self.process_id = os.getpid()
        self.thread_id = threading.get_ident()
        logger.info(f"初始化 WebSearch MCP 服务器")
        
        self.mcp = FastMCP(name)
        self.output_dir = output_dir
        self.headless = headless
        self.format = format
        self.extract_mode = extract_mode
        self.file = file
        
        # 注册工具和资源
        self._register_tools()
        self._register_resources()
        self._register_prompts()
        
        logger.info(f"WebSearch MCP 服务器初始化完成")
    
    def _register_tools(self):
        """注册 MCP 工具"""
        
        @self.mcp.tool()
        def web_search(url: str, format: str = None) -> Dict[str, Any]:
            """
            搜索指定URL并提取网页内容
            
            参数:
                url: 要搜索的网页URL
                format: 返回格式，支持"markdown"、"text"或"html"，默认使用服务器配置的格式
                
            返回:
                包含提取内容和元数据的字典
            """
            logger.info(f"执行 web_search 工具: {url}，格式: {format if format else self.format}")
            
            # 如果指定了format参数，使用指定的格式，否则使用服务器默认配置
            output_format = format if format else self.format
            
            search_tool = WebSearchTool(
                headless=self.headless,
                output_dir=self.output_dir,
                output_format=output_format,
                extract_mode=self.extract_mode
            )
            result = search_tool.search(url)
            return result
        
        @self.mcp.tool()
        def click_element(url: str, element: str, format: str = None) -> Dict[str, Any]:
            """
            访问URL并点击指定元素后提取内容
            
            参数:
                url: 要搜索的网页URL
                element: 要点击的元素选择器
                format: 返回格式，支持"markdown"、"text"或"html"，默认使用服务器配置的格式
                
            返回:
                点击后页面的内容和元数据
            """
            logger.info(f"执行 click_element 工具: {url}, 元素: {element}，格式: {format if format else self.format}")
            
            # 如果指定了format参数，使用指定的格式，否则使用服务器默认配置
            output_format = format if format else self.format
            
            search_tool = WebSearchTool(
                headless=self.headless,
                output_dir=self.output_dir,
                output_format=output_format,
                extract_mode=self.extract_mode
            )
            
            # 访问URL
            search_tool.browser.get(url)
            
            # 点击元素
            search_tool.browser.click(element)
            
            # 等待页面加载
            search_tool.browser.simulate_human_reading(min_time=3.0, max_time=5.0)
            
            # 获取HTML并提取内容
            html_content = search_tool.browser.page.html
            result = search_tool.extractor.extract_content(html_content, url)
            
            # 保存结果
            search_tool._save_result(result)
            
            return result
            
        @self.mcp.tool()
        def search_urls_from_file(format: str = None) -> Dict[str, Any]:
            """
            从文件中读取URL列表并搜索所有URL
            
            参数:
                format: 返回格式，支持"markdown"、"text"或"html"，默认使用服务器配置的格式
                
            返回:
                搜索结果列表
            """
            if not self.file:
                return {"error": "未指定URL文件"}
                
            try:
                with open(self.file, 'r', encoding='utf-8') as f:
                    urls = [line.strip() for line in f if line.strip() and line.strip().startswith(('http://', 'https://'))]
                
                if not urls:
                    return {"error": "文件中未找到有效URL"}
                    
                # 如果指定了format参数，使用指定的格式，否则使用服务器默认配置
                output_format = format if format else self.format
                logger.info(f"从文件 {self.file} 加载了 {len(urls)} 个URL，格式: {output_format}")
                
                results = []
                search_tool = WebSearchTool(
                    headless=self.headless,
                    output_dir=self.output_dir,
                    output_format=output_format,
                    extract_mode=self.extract_mode
                )
                
                for url in urls:
                    try:
                        logger.info(f"处理URL: {url}")
                        result = search_tool.search(url)
                        results.append({
                            "url": url,
                            "success": True,
                            "result": result
                        })
                    except Exception as e:
                        logger.error(f"处理URL时出错 {url}: {str(e)}")
                        results.append({
                            "url": url,
                            "success": False,
                            "error": str(e)
                        })
                
                return {
                    "total": len(urls),
                    "success": sum(1 for r in results if r["success"]),
                    "failed": sum(1 for r in results if not r["success"]),
                    "results": results
                }
                
            except Exception as e:
                logger.error(f"读取URL文件时出错: {str(e)}")
                return {"error": f"读取URL文件时出错: {str(e)}"}
    
    def _register_resources(self):
        """注册 MCP 资源"""
        
        @self.mcp.resource("file://{path}")
        def read_file(path: str) -> str:
            """
            读取指定路径的文件内容
            
            参数:
                path: 文件路径
                
            返回:
                文件内容
            """
            logger.info(f"读取文件资源: {path}")
            try:
                with open(path, 'r', encoding='utf-8') as f:
                    content = f.read()
                return content
            except Exception as e:
                logger.error(f"读取文件失败: {str(e)}")
                return f"错误: {str(e)}"
    
    def _register_prompts(self):
        """注册 MCP 提示模板"""
        
        @self.mcp.prompt()
        def search_web(url: str) -> str:
            """
            创建网页搜索提示模板
            
            参数:
                url: 要搜索的网页URL
                
            返回:
                搜索提示
            """
            return f"请分析以下网页内容，提取关键信息和主要观点: {url}"
    
    def run(self):
        """
        运行 MCP 服务器
        """
        logger.info(f"启动 WebSearch MCP 服务器")
        
        # 直接使用 FastMCP 的 run 方法
        self.mcp.run()
    
    def run_stdio(self):
        """使用标准输入输出运行 MCP 服务器"""
        logger.info("使用标准输入输出启动 WebSearch MCP 服务器")
        import mcp.server.stdio
        import asyncio
        
        async def run_stdio_server():
            async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
                await self.mcp.run_stdio(read_stream, write_stream)
        
        asyncio.run(run_stdio_server()) 